#ifndef UNIASTNODE_H
#define UNIASTNODE_H

#include "antlr4-runtime.h"
#include <string>
namespace UniAST
{
class TreeNode
{
  protected:
    antlr4::ParserRuleContext *ctx;
    std::string nodeName;
    std::vector<TreeNode *> children;
    std::map<antlr4::ParserRuleContext *,
             TreeNode *>
        childrenCtxMap;
    std::map<std::string, std::vector<TreeNode *>>
        childrenNameMap;

  public:
    TreeNode(antlr4::ParserRuleContext *ctx,
             std::string nodeName);

    virtual ~TreeNode();

    void addChild(TreeNode *);
    void deleteChild(TreeNode *);
    TreeNode *
    findChild(antlr4::ParserRuleContext *);

    TreeNode *findChild(std::string name);
    std::vector<TreeNode *>
    findChildren(std::string name);
    virtual bool isTerminal() = 0;
    std::string toStringTree();
    std::string getNodeName();
    antlr4::ParserRuleContext *getCtx();
    std::vector<TreeNode *> &getChildren();
};

class RuleNode : public TreeNode
{
  public:
    RuleNode(antlr4::ParserRuleContext *ctx,
             std::string nodeName);

    bool isTerminal();
};

class TerminalNode : public TreeNode
{
  public:
    TerminalNode(antlr4::ParserRuleContext *ctx,
                 std::string nodeName);

    bool isTerminal();
};
} // namespace UniAST
#endif